from causalnex.structure.notears import from_pandas
from causalnex.network import BayesianNetwork

def train_network(data_val,threshold=0.8):
    net = from_pandas(data_val)
    net.remove_edges_below_threshold(threshold)
    net.get_largest_subgraph()
    return net.get_largest_subgraph()

def train_bn(data_discretised,net):
    bn = BayesianNetwork(net)
    bn = bn.fit_node_states_and_cpds(data = data_discretised,method = "BayesianEstimator",bayes_prior = "K2")
    return bn

